"""
流域边界获取使用示例
基于 mghydro 项目的流域划分功能，通过指定出口点坐标来获取准确的流域边界

参考: https://github.com/mheberger/delineator
"""

import os
import sys
from loguru import logger
import traceback

# 设置环境变量禁用 PyGEOS，强制使用 Shapely 2.0
os.environ['USE_PYGEOS'] = '0'

# 添加项目根目录到 Python 路径
project_root = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.append(project_root)

# 切换到项目根目录，确保能找到 outlets_sample.csv 文件
os.chdir(project_root)

from test_watershed_delineation import WatershedDelineator


def simple_example():
    """简单的流域边界获取示例"""
    logger.info("=== 简单流域边界获取示例 ===")
    
    # 创建流域划分器
    delineator = WatershedDelineator()
    
    try:
        # 定义出口点坐标（冰岛 Nordhura River at Stekkur）
        outlet_id = "6401070"
        latitude = 64.71072
        longitude = -21.60337
        outlet_name = "Nordhura River at Stekkur"
        area_estimate = 507  # 预估面积 507 km²
        
        logger.info(f"出口点ID: {outlet_id}")
        logger.info(f"出口点坐标: ({latitude}, {longitude})")
        logger.info(f"出口点名称: {outlet_name}")
        logger.info(f"预估面积: {area_estimate} km²")
        
        # 获取流域边界
        logger.info("正在计算流域边界...")
        watershed = delineator.delineate_watershed(
            lat=latitude,
            lng=longitude,
            outlet_name=outlet_name,
            outlet_id=outlet_id,
            high_res=True,
            area_estimate=area_estimate
        )
        
        if watershed is not None and len(watershed) > 0:
            logger.success("✓ 流域边界计算成功!")
            logger.info(f"  - 流域数量: {len(watershed)}")
            
            # 计算总面积（平方度）
            total_area = watershed.geometry.area.sum()
            logger.info(f"  - 总面积: {total_area:.6f} 平方度")
            
            # 保存结果为 GeoJSON
            output_file = os.path.join(delineator.output_dir, "simple_watershed.geojson")
            delineator.save_watershed(watershed, output_file, format='geojson')
            logger.success(f"  - 结果已保存到: {output_file}")
            
            return watershed
        else:
            logger.error("✗ 流域边界计算失败")
            return None
            
    except Exception as e:
        traceback.print_exc()
        logger.error(f"✗ 发生错误: {str(e)}")
        return None
    finally:
        delineator.cleanup()


def custom_coordinates_example():
    """自定义坐标的流域边界获取示例"""
    logger.info("=== 自定义坐标流域边界获取示例 ===")
    
    # 创建流域划分器
    delineator = WatershedDelineator()
    
    try:
        # 定义出口点坐标（冰岛 Nordhura River at Stekkur）
        outlet_id = "6401070"
        latitude = 64.71072
        longitude = -21.60337
        outlet_name = "Nordhura River at Stekkur"
        area_estimate = 507  # 预估面积 507 km²
        
        logger.info(f"出口点ID: {outlet_id}")
        logger.info(f"出口点坐标: ({latitude}, {longitude})")
        logger.info(f"出口点名称: {outlet_name}")
        logger.info(f"预估面积: {area_estimate} km²")
        
        # 获取流域边界
        logger.info("正在计算流域边界...")
        watershed = delineator.delineate_watershed(
            lat=latitude,
            lng=longitude,
            outlet_name=outlet_name,
            outlet_id=outlet_id,
            high_res=True,
            area_estimate=area_estimate
        )
        
        if watershed is not None and len(watershed) > 0:
            logger.success("✓ 自定义流域边界计算成功!")
            
            # 保存结果为 GeoJSON
            output_file = os.path.join(delineator.output_dir, "custom_watershed.geojson")
            delineator.save_watershed(watershed, output_file, format='geojson')
            logger.success(f"  - 结果已保存到: {output_file}")
            
            return watershed
        else:
            traceback.print_exc()
            logger.error("✗ 自定义流域计算失败")
            return None
            
    except Exception as e:
        logger.error(f"✗ 自定义流域计算失败: {str(e)}")
        return None
    finally:
        delineator.cleanup()


if __name__ == "__main__":
    logger.info("流域边界获取工具使用示例")
    logger.info("基于 MERIT-Hydro 数据的高精度流域划分")
    logger.info("测试地点：冰岛 Nordhura River at Stekkur (-21.60337, 64.71072)")
    logger.info("=" * 50)
    
    # 运行简单示例
    simple_example()
    
    # 运行自定义坐标示例
    custom_coordinates_example()
    
    logger.info("=" * 50)
    logger.info("使用说明:")
    logger.info("1. 确保已下载必要的 MERIT-Hydro 数据文件")
    logger.info("2. 修改 config.py 中的数据路径设置")
    logger.info("3. 在 custom_coordinates_example() 中修改坐标")
    logger.info("4. 运行此脚本获取流域边界")
    logger.info("数据下载地址:")
    logger.info("- https://mghydro.com/watersheds/")
    logger.info("- https://www.reachhydro.org/home/params/merit-basins")